import triton
import triton.language as tl
import torch
import torch_npu
import pytest
import numpy as np
import random
import test_common
import fa

DEVICE = "npu"
SEED = 42

BATCH_SIZES = [1, 2, 4]
NUM_HEADS = [4, 8, 16]
Q_LENS = [1, None]
KV_LENS = [128, 256, 512, 1024, 2048]
HEAD_DIMS = [32, 64, 128]
DTYPES = [torch.float16, torch.bfloat16]

def naive_fa(
    q, k, v
):
    dtype = q.dtype
    batch_size, num_q_heads, q_len, head_dim = q.shape
    _, num_kv_heads, kv_len, _ = k.shape

    if num_q_heads != num_kv_heads:
        k = torch.repeat_interleave(k, num_q_heads // num_kv_heads, dim=1)
        v = torch.repeat_interleave(v, num_q_heads // num_kv_heads, dim=1)
    
    attn_weight = torch.einsum("bhmd,bhnd->bhmn", q.to(torch.float32), k.to(torch.float32)).to(torch.float32)
    attn_weight = attn_weight.softmax(-1)
    output = torch.einsum("bhmn,bhnd->bhmd", attn_weight, v.to(torch.float32)).to(dtype)
    return output

@pytest.mark.parametrize("batch_size", BATCH_SIZES)
@pytest.mark.parametrize("num_head", NUM_HEADS)
@pytest.mark.parametrize("q_len", Q_LENS)
@pytest.mark.parametrize("kv_len", KV_LENS)
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_fa(
    batch_size,
    num_head,
    q_len,
    kv_len,
    head_dim,
    dtype
):
    Q_LEN = kv_len if q_len is None else q_len
    torch.set_default_device(DEVICE)
    torch.npu.set_device(1)
    torch.manual_seed(SEED)
    np.random.seed(SEED)
    random.seed(SEED)
    sigtype_list = {
        torch.float16: "float16",
        torch.bfloat16: "bfloat16"
    }
    sigtype = sigtype_list[dtype]
    q = test_common.generate_tensor((batch_size, num_head, Q_LEN, head_dim), sigtype).npu()
    k = test_common.generate_tensor((batch_size, num_head, kv_len, head_dim), sigtype).npu()
    v = test_common.generate_tensor((batch_size, num_head, kv_len, head_dim), sigtype).npu()
    
    ref_output = naive_fa(q, k, v)

    actual_out = fa.flash_attn(q, k, v)

    test_common.validate_cmp(sigtype, actual_out, ref_output)


